#include "color.h"
#include "util.h"
#include "frame.h"
#include "table.h"
#include "assem.h"
#include "temp.h"
#include "flowgraph.h"
#include "liveness.h"

U_set precolored                = NULL;
static U_set initial            = NULL;
static U_set simplifyWorklist   = NULL;
static U_set freezeWorklist     = NULL;
static U_set spillWorklist      = NULL;
static U_set spilledWorklist    = NULL;
static U_set coalescedNodes     = NULL;
static U_set coloredNodes       = NULL;
static U_set spilledNodes       = NULL;
static U_setstack selectStack   = NULL;

static U_set coalescedMoves     = NULL;
static U_set constrainedMoves   = NULL;
static U_set frozenMoves        = NULL;
static U_set activeMoves        = NULL;

U_set worklistMoves             = NULL;

U_edgeSet   adjSet              = NULL;
U_nbList    adjList             = NULL;
G_table     worklistN           = NULL;

TAB_table   alias;          // variable A has a alias B if it's combined with B
TAB_table   color;          // color assigned for variables
int         degree[128];    // degree of the nodes

U_set   Adjacent(G_node n);
void    AddEdge(G_node u, G_node v);
void    AddWorkList(G_node u);
void    DecrementDegree(G_node node);
void    EnableMoves(U_set s);
void    MakeWorklist(G_graph graph);
bool    MoveRelated(G_node n);
U_set   NodeMoves(G_node n);
bool    OK(G_node t, G_node r);
void    Simplify();

extern void insertSet(U_set set, Temp_tempList list);

Live_moveList Live_MoveList(G_node src, G_node dst, Live_moveList next)
{
    Live_moveList p = checked_malloc(sizeof(*p));
    p->src = src;
    p->dst = dst;
    p->next = next;
    return p;
}

static inout lookupLiveMap(G_table t,
        G_node flownode)
{
    return (inout) G_look(t, flownode);
}

extern Temp_tempList templist;
extern G_table node_ins;

// map temp to node of conflict map
TAB_table temp_node;

struct COL_conflictGraph BuildConflict(struct Live_graph liveness)
{
    /* G_nodeList nodes    = liveness.nodes; */
    /* TAB_table node_set  = liveness.node_set; */
    /* G_nodeList ptr      = nodes; */
    /* struct COL_conflictGraph lg; */
    /* AS_instr instr      = NULL; */
    /* Temp_tempList tempptr = templist;     //list of all temp vars */
    /* U_setIter iter, iter2; */
    /* U_set def, live; */

    /* worklistN   = TAB_empty(); */
    /* temp_node = TAB_empty(); */
    /* lg.graph = G_Graph(); */
    /* worklistMoves = U_setNew(); */
    
    /* ptr = nodes; */
    /* while(ptr->tail) */
    /*     ptr = ptr->tail; */
    /* live = lookupLiveMap(node_set, ptr->head)->out; */


    /* while(tempptr)             //add all temp vars into conflict graph */
    /* { */
    /*     G_node node; */

    /*     assert(tempptr->head); */
    /*     node = G_Node(lg.graph, tempptr->head); */
    /*     TAB_enter(temp_node, tempptr->head, node); */
    /*     tempptr = tempptr->tail; */
    /* } */

    /* // loop for all instruction nodes in liveness gragh */
    /* lg.moves = NULL; */
    /* while(ptr) */
    /* { */
    /*     instr = TAB_look(node_ins, ptr->head); */

    /*     if (FG_isMove(ptr->head))    // if it is a MOVE instruction */
    /*     { */
            /* U_set use = U_setNew(); */
            /* Temp_temp src = NULL, des = NULL; */
            /* Temp_tempList ptrTemp; */
            /* U_set setMovelistN; */
            /* setMovelistN = U_setNew(); */

            /* insertSet(use, FG_use(ptr->head)); */
            /* U_setSub2(live, use); */

            /* // find out the source and destination of the move instruction */
            /* if (instr->u.MOVE.src) */
            /*     src = instr->u.MOVE.src->head; */
            /* if (instr->u.MOVE.dst) */
            /*     des = instr->u.MOVE.dst->head; */

            /* // worklistN: G_node => U_set( temp var relative to MOVE ) */
            /* for(ptrTemp = instr->u.MOVE.src; ptrTemp; ptrTemp = ptrTemp->tail) */
            /*     U_setInsert(setMovelistN, ptrTemp->head); */
            /* for(ptrTemp = instr->u.MOVE.dst; ptrTemp; ptrTemp = ptrTemp->tail) */
            /*     U_setInsert(setMovelistN, ptrTemp->head); */
            /* G_enter(worklistN, ptr->head, setMovelistN); */

            /* // add the MOV instruction into worklistMoves */
            /* U_setInsert(worklistMoves, ptr->head); */

            /* lg.moves = Live_MoveList( */
            /*         (G_node)TAB_look(temp_node, src), */
            /*         (G_node)TAB_look(temp_node, des), */
            /*         lg.moves); */
 
/*             // TODO: What does this supposed to do? */
/*             //iter = U_setFirst(io->out); */
/*             //while(iter) */
/*             //{ */
/*             //    // MOV a, b */
/*             //    // add edge to the conflict graph */
/*             //    Temp_temp tempc = U_setLookup(iter); */
/*             //    Temp_temp tempb; */
/*             //    Temp_temp tempa; */
/*             //    if (instr->u.MOVE.src) */
/*             //        tempb = instr->u.MOVE.src->head; */
/*             //    else */
/*             //        tempb = NULL; */
/*             //    if (instr->u.MOVE.dst) */
/*             //        tempa = instr->u.MOVE.dst->head; */
/*             //    else */
/*             //        tempa = NULL; */

/*             //    if (tempb && tempb != tempc) */
/*             //    { */
/*             //        G_node a = TAB_look(temp_node, tempa); */
/*             //        G_node b = TAB_look(temp_node, tempb); */

/*             //        G_addEdge(a, b); */
/*             //    } */
/*             //    U_setWalk(&iter); */
/*             //} */
/* // U_setDestroy(use); */
/*         } // eof if move instruction */

/*         def = U_setNew(); */
/*         insertSet(def, FG_def(ptr->head)); */
/*         U_setAdd2(live, def); */

/*         iter = U_setFirst(live); */
/*         iter2 = U_setFirst(def); */
/*         while(iter2) */
/*         { */
/*             while(iter) */
/*             { */
/*                 Temp_temp tempb = U_setLookup(iter); */
/*                 Temp_temp tempa = U_setLookup(iter2); */
/*                 G_node b = TAB_look(temp_node, tempb); */
/*                 G_node a = TAB_look(temp_node, tempa); */

/*                 AddEdge(a, b); */
/*                 U_setWalk(&iter); */
/*             } */
/*             U_setWalk(&iter); */
/*         } */

/*         U_setDestroy(def); */
/*         ptr = ptr->tail; */
    /* } */

    /* return lg; */
};

U_set NodeMoves(G_node n)
{
    /* U_set temp = U_setNew(); */
    /* U_set result = U_setNew(); */
    /* U_set movelist; */

    /* assert(n); */
    /* movelist = G_look(worklistN, n); */
    /* if(movelist) */
    /* { */
    /*     U_setAdd3(temp, activeMoves, worklistMoves); */
    /*     U_setSub3(result, movelist, temp); */
    /* } */
    /* U_setDestroy(temp); */
    /* return result; */
}

U_set Adjacent(G_node n)
{
    U_set result = U_setNew();
    U_set temp = U_setNew();
    int i, key;

    U_setAdd3(temp, U_setstackCloneSet(selectStack), coalescedNodes);
    key = U_nbListGetKey(adjList, n);
    for(i=0; i< U_nbListCount(adjList); i++)
        if ( U_nbListIsAdjK(adjList, key, i))
        {
            G_node adjnode = U_nbListGetPtr(adjList, i);
            if (!U_setHas(temp, adjnode))
                U_setInsert(result, adjnode);
        }
    return result;
}

bool MoveRelated(G_node n)
{
    U_set set;
    int count;

    assert(n);
    set = NodeMoves(n);
    count = U_setGetCount(set);
    U_setDestroy(&set);
    return (bool) (count != 0);
}

void MakeWorklist(G_graph graph)
{
    G_nodeList ptr1, ptr2;
    U_set initial = U_setNew();
    U_setIter iter = NULL;
    int K = F_getRegNum();
    int i;

    assert(graph);
    ptr1 = ptr2 = G_nodes(graph);
    spillWorklist    = U_setNew();
    freezeWorklist    = U_setNew();
    simplifyWorklist= U_setNew();
    adjList = U_nbListNew();

    while(ptr1)
    {
        U_setInsert(initial, ptr1->head);
        U_nbListEnter(adjList, ptr1->head);
        ptr1 = ptr1->tail;
    }

    // create the adjacent list and degree array
    ptr1 = ptr2;
    while(ptr1)
    {
        int key;
        while(ptr2)
        {
            if (G_goesTo(ptr1->head, ptr2->head))
                U_nbListAddEdge(adjList, ptr1->head, ptr2->head);
            ptr2 = ptr2->tail;
        }
        key = U_nbListGetKey(adjList, ptr1);
        degree[key] = G_degree(ptr1->head);
        ptr1 = ptr1->tail;
    }

    for(; U_setCount(initial);)
    {
        G_node n = (G_node)U_setLookup(initial, 0);
        if (G_degree(n) > K)
            U_setInsert(spillWorklist, n);
        else if (MoveRelated(n))
            U_setInsert(freezeWorklist, n);
        else
            U_setInsert(simplifyWorklist, n);

        U_setRemove(initial, n);
    }
    U_setDestroy(&initial);
}

void EnableMoves(U_set s)
{
    /* U_setIter iter1 = U_setFirst(s); */
    /* while(iter1) */
    /* { */
    /*     G_node n = U_setLookup(iter1); */
    /*     U_set tempset = NodeMoves(n); */
    /*     U_setIter iter2 = U_setFirst(tempset); */
    /*     while(iter2) */
    /*     { */
    /*         G_node data = U_setLookup(iter2); */
    /*         if ( U_setHas(activeMoves, data ) ) */
    /*         { */
    /*             U_setRemove(activeMoves, data); */
    /*             U_setInsert(worklistMoves, data); */
    /*         } */
    /*         U_setWalk(&iter2); */
    /*     } */
    /*     U_setDestroy(tempset); */

    /*     U_setWalk(&iter1); */
    /* } */
}

void DecrementDegree(G_node node)
{
    int key = U_nbListGetKey(adjList, node);
    int d = degree[key] --;
    if ( d == F_getRegNum() )
    {
        U_set tempset = Adjacent(node);
        U_setInsert(tempset, node);
        EnableMoves(tempset);
        U_setRemove(spillWorklist, node);
        if (MoveRelated(node))
            U_setRemove(freezeWorklist, node);
        else
            U_setRemove(simplifyWorklist, node);
    }
}

void Simplify()
{
    /* U_setIter iter1; */

    /* iter1 = U_setFirst(simplifyWorklist); */
    /* while(iter1) */
    /* { */
    /*     U_setIter iter2; */
    /*     G_node n = U_setLookup(iter1); */
    /*     U_set tempset; */

    /*     U_setRemove(simplifyWorklist, n); */
    /*     U_setstackPush(selectStack, n); */
        
    /*     tempset = Adjacent(n); */
    /*     iter2 = U_setFirst(tempset); */
    /*     while(iter2) */
    /*     { */
    /*         DecrementDegree( U_setLookup(iter2) ); */
    /*         U_setWalk(&iter2); */
    /*     } */

    /*     U_setWalk(&iter1); */
    /* } */
}

void AddWorkList(G_node u)
{
    if (!U_setHas(precolored, u)
        && !MoveRelated(u)
        && degree[ U_nbListGetKey(adjList, u) ] < F_getRegNum() )
    {
        U_setRemove(freezeWorklist, u);
        U_setInsert(simplifyWorklist, u);
    }
}

bool OK(G_node t, G_node r)
{
    return (bool) ( degree[ U_nbListGetKey(adjList, t) ] < F_getRegNum()
        || U_setHas(precolored, t)
        || U_edgeSetIsAdj(adjSet, t, r) );
}

bool Conservative(U_set nodes)
{
    /* int k=0; */
    /* U_setIter iter = U_setFirst(nodes); */
    /* while(iter) */
    /* { */
    /*     G_node n = U_setLookup(iter); */
    /*     if (degree[ U_nbListGetKey(adjList, n) ] >= F_getRegNum() ) */
    /*         k++; */
    /*     U_setWalk(&iter); */
    /* } */
    /* return (bool) ( k<F_getRegNum() ); */
}

void Combine(G_node u, G_node v)
{
    /* U_set moveListU, moveListV; */
    /* U_set adjV; */
    /* U_setIter adjIter; */

    /* if ( U_setHas(freezeWorklist, v) ) */
    /*     U_setRemove(freezeWorklist, v); */
    /* else */
    /*     U_setRemove(spillWorklist, v); */
    /* U_setInsert(coalescedNodes, v); */
    /* TAB_enter(alias, v, u); */

    /* moveListU = G_look(worklistN, u); */
    /* moveListV = G_look(worklistN, v); */
    /* U_setAdd2(moveListU, moveListV); */
    /* EnableMoves(moveListV); */

    /* adjV = Adjacent(v); */
    /* adjIter = U_setFirst(adjV); */
    /* while(adjIter) */
    /* { */
    /*     G_node t = U_setLookup(adjIter); */
    /*     AddEdge(t, u); */
    /*     DecrementDegree(t); */
    /*     U_setWalk(&adjIter); */
    /* } */
    /* U_setDestroy(adjV); */

    /* if ( degree[ U_nbListGetKey(adjList, u) ] >= F_getRegNum() */
    /*     && U_setHas(freezeWorklist, u) ) */
    /* { */
    /*     U_setRemove(freezeWorklist, u); */
    /*     U_setInsert(spillWorklist, u); */
    /* } */
}

void Coalesce()
{
    /* U_setIter iter1 = U_setFirst(worklistMoves); */

    /* while(iter1) */
    /* { */
    /*     G_node m = U_setLookup(iter1); */
    /*     AS_instr instr = FG_getInstr(m); */
    /*     Temp_temp x = instr->u.MOVE.dst->head; */
    /*     Temp_temp y = instr->u.MOVE.src->head; */
    /*     G_node u,v; */

    /*     if (U_setHas(precolored, y)) */
    /*     { */
    /*         Temp_temp t = x; */
    /*         x = y; */
    /*         y = t; */
    /*     } */
    /*     U_setRemove(worklistMoves, m); */
    /*     u = Live_tempToGnode(x); */
    /*     v = Live_tempToGnode(y); */

    /*     if (u == v) */
    /*     { */
    /*         U_setInsert(coalescedMoves, m); */
    /*         AddWorkList(u); */
    /*         U_setWalk(&iter1); */
    /*         continue; */
    /*     } */
    /*     else if ( U_setHas(precolored, v) ) */
    /*     { */
    /*         if (U_edgeSetIsAdj(adjSet, u, v)) */
    /*         { */
    /*             U_setInsert(constrainedMoves, m); */
    /*             AddWorkList(u); */
    /*             AddWorkList(v); */
    /*             U_setWalk(&iter1); */
    /*             continue; */
    /*         } */
    /*         else */
    /*         { */
    /*             U_set vadj = Adjacent(v); */
    /*             U_setIter viter = U_setFirst(vadj); */
    /*             while( viter) */
    /*             { */
    /*                 G_node t = U_setLookup(viter); */
    /*                 if( OK(t, u) ) */
    /*                 { */
    /*                     U_setInsert(coalescedMoves, m); */
    /*                     Combine(u, v); */
    /*                     AddWorkList(u); */
    /*                     U_setDestroy(vadj); */
    /*                     U_setWalk(&iter1); */
    /*                     continue; */
    /*                 } */
    /*                 U_setWalk(&viter); */
    /*             } */
    /*         } */
    /*     } */
    /*     else */
    /*     { */
    /*         U_set uadj = Adjacent(u); */
    /*         U_set vadj = Adjacent(v); */
    /*         U_set temp = U_setNew(); */
    /*         U_setAdd3(temp, uadj, vadj); */
    /*         U_setDestroy(uadj); U_setDestroy(vadj); */

    /*         if (Conservative(temp)) */
    /*         { */
    /*             U_setInsert(coalescedMoves, m); */
    /*             Combine(u, v); */
    /*             AddWorkList(u); */
    /*             U_setWalk(&iter1); */
    /*             continue; */
    /*         } */
    /*         U_setDestroy(temp); */
    /*     } */

    /*     U_setInsert(activeMoves, m); */
    /*     U_setWalk(&iter1); */
    /* } // eof while(iter1) */
}

G_node GetAlias(G_node n)
{
    if (U_setHas(coalescedNodes, n))
        return GetAlias((G_node)(TAB_look(alias, n)));
    return n;
}

void FreezeMoves(G_node u)
{
    /* U_set nodemoves = NodeMoves(u); */
    /* U_setIter iter = U_setFirst(nodemoves); */
    
    /* while(iter) */
    /* { */
    /*     G_node m = U_setLookup(iter); */
    /*     AS_instr instr = FG_getInstr(m); */
    /*     Temp_temp a = instr->u.MOVE.dst->head; */
    /*     Temp_temp b = instr->u.MOVE.src->head; */
    /*     G_node x, y, v; */
    /*     U_set temp; */

    /*     x = Live_tempToGnode(a); */
    /*     y = Live_tempToGnode(b); */

    /*     if ( GetAlias(y) == GetAlias(u) ) */
    /*         v = GetAlias(x); */
    /*     else */
    /*         v = GetAlias(y); */

    /*     U_setRemove(activeMoves, m); */
    /*     U_setInsert(frozenMoves, m); */
    /*     temp = NodeMoves(v); */
    /*     if ( U_setGetCount(temp) == 0 */
    /*         && degree[ U_nbListGetKey(adjList, v) ] < F_getRegNum()) */
    /*     { */
    /*         U_setRemove(freezeWorklist, v); */
    /*         U_setInsert(simplifyWorklist, v); */
    /*     } */
    /*     U_setWalk(&iter); */
    /* } */
}

void Freeze()
{
    /* U_setIter iter = U_setFirst(freezeWorklist); */
    /* while(iter) */
    /* { */
    /*     G_node u = U_setLookup(iter); */
    /*     U_setRemove(freezeWorklist, u); */
    /*     U_setInsert(simplifyWorklist, u); */
    /*     FreezeMoves(u); */
    /*     U_setWalk(&iter); */
    /* } */
}

void SelectSpill()
{
    /* // TODO: change the way we select node from spillWorklist to avoid selecting */
    /* // specific nodes. */
    /* U_setIter iter = U_setFirst(spillWorklist); */
    /* while(iter) */
    /* { */
    /*     G_node m = U_setLookup(iter); */
    /*     U_setRemove(spillWorklist, m); */
    /*     U_setInsert(simplifyWorklist, m); */
    /*     FreezeMoves(m); */
    /*     U_setWalk(&iter); */
    /* } */
}

void AssignColors()
{    
    /* U_setIter iter; */
    /* while( U_setstackCount(selectStack) ) */
    /* { */
    /*     U_set okColors = U_setNew(); */
    /*     int i, key; */
    /*     G_node n; */

    /*     n = U_setstackPop(selectStack); */

    /*     for (i=0; i<F_getRegNum(); i++) */
    /*         U_setInsert(okColors, (void*)i); */

    /*     key = U_nbListGetKey(adjList, n); */
    /*     for(i=0; i< U_nbListCount(adjList); i++) */
    /*         if ( U_nbListIsAdjK(adjList, key, i)) */
    /*         { */
    /*             G_node w = (G_node) U_nbListGetPtr(adjList, i); */
    /*             if (U_setHas(coloredNodes, GetAlias(w) ) */
    /*                 || U_setHas(precolored, GetAlias(w) )) */
    /*             { */
    /*                 int c = (int) TAB_look(color, GetAlias(w)); */
    /*                 U_setRemove(okColors, (void*)c); */
    /*             } */
    /*         } */
    /*     if( U_setGetCount(okColors) ) */
    /*     { */
    /*         void* c; */
    /*         iter = U_setFirst(okColors); */
    /*         U_setInsert(coloredNodes, n); */
    /*         c = U_setLookup(iter); */
    /*         TAB_enter(color, n, c); */
    /*     } */
    /*     else */
    /*     { */
    /*         U_setInsert(spilledNodes, n); */
    /*     } */

    /*     U_setDestroy(okColors); */
    /* } */

    /* for(iter = U_setFirst(coalescedNodes); */
    /*     iter; */
    /*     U_setWalk(&iter) ) */
    /* { */
    /*     G_node n = (G_node) U_setLookup(iter); */
    /*     TAB_enter(color, n, TAB_look(color, GetAlias(n) ) ); */
    /* } */
}

void RewriteProgram()
{


}

void COL_init()
{
    selectStack     = U_setstackNew();
    coalescedNodes  = U_setNew();
    coloredNodes    = U_setNew();
    spilledNodes    = U_setNew();
    color           = TAB_empty();
}

struct COL_result COL_colorVar(G_graph flowGraph)
{
    struct COL_result result;
    struct Live_graph liveGraph = Live_liveness(flowGraph);
    struct COL_conflictGraph confGraph = BuildConflict(liveGraph);
    Temp_tempList ptr;

    MakeWorklist(confGraph.graph);

    do
    {
        if (U_setGetCount(simplifyWorklist))
            Simplify();
        else if (U_setGetCount(worklistMoves))
            Coalesce();
        else if (U_setGetCount(freezeWorklist))
            Freeze();
        else if (U_setGetCount(spillWorklist))
            SelectSpill();
    }while( !U_setGetCount(simplifyWorklist)
            && !U_setGetCount(worklistMoves)
            && !U_setGetCount(freezeWorklist)
            && !U_setGetCount(spillWorklist));
    AssignColors();

    // TODO
    if( U_setGetCount(spilledNodes) )
    {
        assert( 0 );
        //rewriteProgram(spilledNodes);
        //COL_colorVar();
    }

    result.coloring = Temp_empty();
    result.spills = NULL;
    ptr = templist;
    while(ptr)
    {
        int c;
        char buff[255];
        Temp_temp temp = ptr->head;
        G_node n = TAB_look(temp_node, temp);
        c = (int)TAB_look(color, n);
        
        buff[0] = 'R';
        buff[1] = 'E';
        buff[2] = 'G';
        buff[3] = '0' + c;
        buff[4] = 0;

        assert(n);
        c = (int) TAB_look(color, n);
        Temp_enter(result.coloring, temp, buff);
        
        ptr = ptr->tail;
    }
    return result;
}

void AddEdge(G_node u, G_node v)
{
    int key;
    if( U_edgeSetIsAdj(adjSet, u, v) && u != v )
    {
        U_edgeSetAddEdge(adjSet, u, v);
        U_edgeSetAddEdge(adjSet, v, u);
        if ( !U_setHas( precolored, u) )
        {
            U_nbListAddEdge(adjList, u, v);
            key = U_nbListGetKey(adjList, u);
            degree[key] ++;
        }
        if ( !U_setHas( precolored, v) )
        {
            U_nbListAddEdge(adjList, v, u);
            key = U_nbListGetKey(adjList, v);
            degree[key] ++;
        }
    }
}
